package top.lingkang.finalgateway.core.tcp;

import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import top.lingkang.finalgateway.config.TcpConfig;

/**
 * @author lingkang
 * @create by 2024/5/23 12:15
 */
public class TcpProxyHandler extends ChannelInboundHandlerAdapter {
    private final TcpConfig config;
    private Channel outboundChannel;
    private long total = 0;
    private final String ip;

    public TcpProxyHandler(TcpConfig config, String ip) {
        this.config = config;
        this.ip = ip;
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        final Channel inboundChannel = ctx.channel();

        // Start the connection attempt.
        Bootstrap b = new Bootstrap();
        b.group(inboundChannel.eventLoop())
                .channel(ctx.channel().getClass())
                .handler(new TcpProxyBackendHandler(inboundChannel))
                .option(ChannelOption.AUTO_READ, false);
        ChannelFuture f = b.connect(config.getTargetHost(), config.getTargetPort());
        outboundChannel = f.channel();
        f.addListener(new ChannelFutureListener() {
            @Override
            public void operationComplete(ChannelFuture future) {
                if (future.isSuccess()) {
                    // 连接完成开始读取第一个数据
                    inboundChannel.read();
                } else {
                    // 如果连接尝试失败，关闭连接。
                    inboundChannel.close();
                }
            }
        });
    }

    @Override
    public void channelRead(final ChannelHandlerContext ctx, Object msg) {
        if (outboundChannel.isActive()) {
            if (msg instanceof ByteBuf) {
                ByteBuf buf = (ByteBuf) msg;
                total += buf.readableBytes();
            }
            outboundChannel.writeAndFlush(msg).addListener(new ChannelFutureListener() {
                @Override
                public void operationComplete(ChannelFuture future) {
                    if (future.isSuccess()) {
                        // 能够清除数据，开始读取下一个区块
                        ctx.channel().read();
                    } else {
                        future.channel().close();
                    }
                }
            });
        }
    }

    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
        config.getMonitor().accessTcp(ip, config.getPort(), config.getTargetHost(), config.getTargetPort(), total);
    }

    static void closeOnFlush(Channel ch) {
        if (ch.isActive()) {
            ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
        }
    }
}
